# pip install opencv-python mediapipe numpy websockets

import cv2
import mediapipe as mp
import numpy as np
import socket
import json
import time
import websockets
from collections import deque
import threading
import asyncio
from typing import Optional

# 传输配置
TRANSPORT = "udp"  # "udp" | "tcp" | "ws"
HOST = "127.0.0.1"
PORT = 8765
WS_URL = "ws://127.0.0.1:8765"
SOURCE_ID = "py-hand-1"  # 事件来源标识

# 视觉与手势配置
MIRROR_PREVIEW = False
PROCESS_LEFT = True
PROCESS_RIGHT = True
DEBUG_PRINT = False
PRINT_JOINTS_EVERY_N = 0

REFERENCE_HAND_WIDTH_M = 0.06
EMA_ALPHA = 0.45

# 抬手参数
RAISE_Y_THRESHOLD_M = 0.04
RAISE_SPREAD_MIN_M = 0.050
RAISE_HOLD_FRAMES = 3
RAISE_COOLDOWN_S = 1.0

# 抓取参数（距离 + 角度）
GRAB_ENTER_AVG_TIP_DIST_M = 0.045
GRAB_EXIT_AVG_TIP_DIST_M  = 0.055
GRAB_MIN_TIPS_NEAR = 2
GRAB_HOLD_FRAMES = 2
GRAB_COOLDOWN_S = 0.5
GRAB_TIP_TO_PALM_COS_MIN = 0.35
GRAB_MIN_TIPS_ANG_OK = 3

# 挥手参数
WAVE_WINDOW_S = 1.2
WAVE_MIN_REVERSALS = 2
WAVE_MIN_AMPLITUDE_M = 0.06
WAVE_COOLDOWN_S = 1.2

# MediaPipe 关键点索引
JOINTS = {
    "Wrist": 0,
    "ThumbTip": 4,
    "IndexMCP": 5,
    "IndexPIP": 6,
    "IndexTip": 8,
    "MiddleMCP": 9,
    "MiddlePIP": 10,
    "MiddleTip": 12,
    "RingMCP": 13,
    "RingPIP": 14,
    "RingTip": 16,
    "PinkyMCP": 17,
    "PinkyPIP": 18,
    "PinkyTip": 20
}

# 传输层封装（带 seq/source，TCP/WS 自动重连）
class Transport:
    def __init__(self, mode, host=None, port=None, ws_url=None, source_id="source"):
        self.mode = mode.lower()
        self.host = host
        self.port = port
        self.ws_url = ws_url
        self.source_id = source_id

        self.seq = 0  # 事件序号
        self.udp_sock: Optional[socket.socket] = None
        self.tcp_sock: Optional[socket.socket] = None

        # WS 异步支持
        self.ws_loop: Optional[asyncio.AbstractEventLoop] = None
        self.ws_thread: Optional[threading.Thread] = None
        self.ws_send_queue: Optional[asyncio.Queue] = None
        self.ws_connected = False
        self.ws_stop = False

    # 通用消息
    def build_data(self, event_dict: dict) -> bytes:
        self.seq = (self.seq + 1) & 0xFFFFFFFF  # 32-bit wrap
        event_dict["seq"] = self.seq
        event_dict["source"] = self.source_id
        return json.dumps(event_dict, ensure_ascii=False).encode("utf-8")

    def start(self):
        if self.mode == "udp":
            self.udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            print(f"[NET] UDP -> {self.host}:{self.port}")
        elif self.mode == "tcp":
            print(f"[NET] TCP connect -> {self.host}:{self.port}")
            self._tcp_connect_with_retry()
        elif self.mode == "ws":
            print(f"[NET] WS connect -> {self.ws_url}")
            self._ws_start_thread()
        else:
            raise ValueError("TRANSPORT must be 'udp' | 'tcp' | 'ws'")

    # UDP
    def _udp_send(self, data: bytes):
        try:
            self.udp_sock.sendto(data, (self.host, self.port))
        except Exception as e:
            print(f"[ERR] UDP send: {e}")

    # TCP
    def _tcp_connect_with_retry(self):
        backoff = 0.5
        max_backoff = 5.0
        while True:
            try:
                self.tcp_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.tcp_sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                self.tcp_sock.settimeout(5.0)
                self.tcp_sock.connect((self.host, self.port))
                self.tcp_sock.settimeout(None)
                print("[NET] TCP connected")
                return
            except Exception as e:
                print(f"[WARN] TCP connect failed: {e}; retry in {backoff:.1f}s")
                time.sleep(backoff)
                backoff = min(max_backoff, backoff * 1.7 + 0.1)

    def _tcp_send(self, data: bytes):
        try:
            header = len(data).to_bytes(4, byteorder="big")
            self.tcp_sock.sendall(header + data)
        except Exception as e:
            print(f"[WARN] TCP send failed: {e}; reconnecting...")
            try:
                self.tcp_sock.close()
            except Exception:
                pass
            self._tcp_connect_with_retry()
            # 再尝试一次发送
            try:
                header = len(data).to_bytes(4, byteorder="big")
                self.tcp_sock.sendall(header + data)
            except Exception as e2:
                print(f"[ERR] TCP resend failed: {e2}")

    # WebSocket
    def _ws_start_thread(self):
        self.ws_stop = False
        self.ws_thread = threading.Thread(target=self._ws_thread_main, daemon=True)
        self.ws_thread.start()

    def _ws_thread_main(self):
        try:
            asyncio.set_event_loop(asyncio.new_event_loop())
        except RuntimeError:
            pass
        self.ws_loop = asyncio.get_event_loop()
        self.ws_send_queue = asyncio.Queue()
        self.ws_loop.run_until_complete(self._ws_run())

    async def _ws_run(self):
        backoff = 0.5
        max_backoff = 5.0
        while not self.ws_stop:
            try:
                async with websockets.connect(self.ws_url, max_queue=None) as ws:
                    print("[NET] WS connected")
                    self.ws_connected = True
                    backoff = 0.5
                    # 发送循环
                    while not self.ws_stop:
                        data = await self.ws_send_queue.get()
                        if data is None:
                            break
                        try:
                            await ws.send(data)
                        except Exception as e:
                            print(f"[WARN] WS send failed: {e}")
                            break
            except Exception as e:
                print(f"[WARN] WS connect failed: {e}; retry in {backoff:.1f}s")
                await asyncio.sleep(backoff)
                backoff = min(max_backoff, backoff * 1.7 + 0.1)
            finally:
                self.ws_connected = False

    def _ws_send(self, data: bytes):
        if not self.ws_loop or not self.ws_send_queue:
            return
        # bytes发送为JSON
        def _put():
            if not self.ws_send_queue:
                return
            try:
                self.ws_send_queue.put_nowait(data.decode("utf-8"))
            except Exception as e:
                print(f"[ERR] WS enqueue: {e}")
        # 投递任务
        self.ws_loop.call_soon_threadsafe(_put)

    def send(self, data: bytes):
        if self.mode == "udp":
            self._udp_send(data)
        elif self.mode == "tcp":
            self._tcp_send(data)
        elif self.mode == "ws":
            self._ws_send(data)

    def close(self):
        try:
            if self.udp_sock:
                self.udp_sock.close()
            if self.tcp_sock:
                self.tcp_sock.close()
            if self.mode == "ws":
                self.ws_stop = True
                if self.ws_send_queue and self.ws_loop:
                    # 结束发送循环
                    def _stop():
                        try:
                            self.ws_send_queue.put_nowait(None)
                        except Exception:
                            pass
                    self.ws_loop.call_soon_threadsafe(_stop)
                if self.ws_thread:
                    self.ws_thread.join(timeout=2.0)
        except Exception:
            pass

# 手势识别实现
mp_hands = mp.solutions.hands
hands = mp_hands.Hands(
    static_image_mode=False,
    max_num_hands=2,
    min_detection_confidence=0.5,
    min_tracking_confidence=0.5,
    model_complexity=1
)

prev_points = {"left": None, "right": None}
last_grab_time = {"left": 0.0, "right": 0.0}
last_raise_time = {"left": 0.0, "right": 0.0}
last_wave_time  = {"left": 0.0, "right": 0.0}
wave_history = {"left": deque(), "right": deque()}  # (t, x_m)
hold_counters = {"left": {"raise": 0, "grab": 0}, "right": {"raise": 0, "grab": 0}}
frame_count = 0

POINT_ORDER = [
    "Wrist",
    "ThumbTip",
    "IndexMCP", "IndexPIP", "IndexTip",
    "MiddleMCP", "MiddlePIP", "MiddleTip",
    "RingMCP", "RingPIP", "RingTip",
    "PinkyMCP", "PinkyPIP", "PinkyTip"
]

def ema(prev, cur, alpha):
    if prev is None:
        return cur
    return prev * (1 - alpha) + cur * alpha

def send_event(net: Transport, hand, gesture):
    payload = {"t": time.time(), "hand": hand, "gesture": gesture}
    data = net.build_data(payload)
    net.send(data)
    print(f"[GESTURE] {hand:<5s} -> {gesture} (seq={net.seq})")

def unit_vec(v):
    n = np.linalg.norm(v)
    if n < 1e-6:
        return v * 0.0
    return v / n

def compute_metrics(points_dict):
    mcp_names = ["IndexMCP", "MiddleMCP", "RingMCP", "PinkyMCP"]
    mcp = np.stack([points_dict[n] for n in mcp_names], axis=0)
    palm = mcp.mean(axis=0)

    tip_names = ["ThumbTip", "IndexTip", "MiddleTip", "RingTip", "PinkyTip"]
    tips = np.stack([points_dict[n] for n in tip_names], axis=0)

    tip_dists = np.linalg.norm(tips - palm[None, :], axis=1)
    spread = float(tip_dists.mean())

    index_dir = points_dict["IndexTip"] - points_dict["IndexMCP"]
    index_dir_y = float(index_dir[1])

    near_mask = (tip_dists < GRAB_ENTER_AVG_TIP_DIST_M)
    near_count = int(near_mask.sum())

    return {
        "palm": palm,
        "spread": spread,
        "grab_avg": float(tip_dists.mean()),
        "raise_y": float(palm[1]),
        "index_dir_y": index_dir_y,
        "index_tip": points_dict["IndexTip"],
        "tip_names": tip_names
    }

def compute_grab_angle_metrics(points_dict, tip_names):
    palm_candidates = ["IndexMCP", "MiddleMCP", "RingMCP", "PinkyMCP"]
    palm = np.stack([points_dict[n] for n in palm_candidates], axis=0).mean(axis=0)

    cos_list = []
    for tn in tip_names:
        if "Index" in tn:
            pip = points_dict["IndexPIP"]
        elif "Middle" in tn:
            pip = points_dict["MiddlePIP"]
        elif "Ring" in tn:
            pip = points_dict["RingPIP"]
        elif "Pinky" in tn:
            pip = points_dict["PinkyPIP"]
        else:
            pip = points_dict["IndexMCP"]  # Thumb 粗略
        tip = points_dict[tn]
        v_tip_dir = unit_vec(tip - pip)
        v_to_palm = unit_vec(palm - tip)
        cos_val = float(np.dot(v_tip_dir, v_to_palm))
        cos_list.append(cos_val)
    cos_arr = np.array(cos_list, dtype=np.float32)
    ang_ok = int(np.sum(cos_arr >= GRAB_TIP_TO_PALM_COS_MIN))
    return ang_ok, cos_arr

def update_wave(hand, t_now, index_tip_x_m):
    dq = wave_history[hand]
    dq.append((t_now, index_tip_x_m))
    while dq and (t_now - dq[0][0]) > WAVE_WINDOW_S:
        dq.popleft()
    if len(dq) < 5:
        return 0, 0.0
    xs = np.array([x for _, x in dq])
    dx = np.diff(xs)
    signs = np.sign(dx)
    signs[abs(dx) < 1e-4] = 0
    reversals = 0
    last = 0
    for s in signs:
        if s == 0:
            continue
        if last == 0:
            last = s
        elif s != last:
            reversals += 1
            last = s
    amplitude = float(xs.max() - xs.min())
    return reversals, amplitude

def main():
    global frame_count

    net = Transport(TRANSPORT, host=HOST, port=PORT, ws_url=WS_URL, source_id=SOURCE_ID)
    net.start()

    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("Failed to open camera.")
        net.close()
        return

    print(f"[INFO] Transport={TRANSPORT} (udp/tcp host={HOST}:{PORT}, ws={WS_URL})")
    print(f"[INFO] MIRROR_PREVIEW={MIRROR_PREVIEW}")

    mp_hands_local = mp.solutions.hands
    hands_local = mp_hands_local.Hands(
        static_image_mode=False,
        max_num_hands=2,
        min_detection_confidence=0.5,
        min_tracking_confidence=0.5,
        model_complexity=1
    )

    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            frame_count += 1

            if MIRROR_PREVIEW:
                frame = cv2.flip(frame, 1)

            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = hands_local.process(rgb)
            h, w = frame.shape[:2]

            if results.multi_hand_landmarks and results.multi_handedness:
                for hand_landmarks, handedness in zip(results.multi_hand_landmarks, results.multi_handedness):
                    label = handedness.classification[0].label.lower()
                    if MIRROR_PREVIEW:
                        label = "left" if label == "right" else "right"
                    if (label == "left" and not PROCESS_LEFT) or (label == "right" and not PROCESS_RIGHT):
                        continue

                    lm = hand_landmarks.landmark

                    p_index_tip_px = np.array([lm[8].x * w, lm[8].y * h])
                    p_pinky_tip_px = np.array([lm[20].x * w, lm[20].y * h])
                    pixel_width = max(np.linalg.norm(p_index_tip_px - p_pinky_tip_px), 1.0)
                    meters_per_pixel = REFERENCE_HAND_WIDTH_M / pixel_width

                    pts = {}
                    for name, idx in JOINTS.items():
                        x_px = lm[idx].x * w
                        y_px = lm[idx].y * h
                        x_c = (x_px - w / 2.0)
                        y_c = -(y_px - h / 2.0)
                        z_px = -lm[idx].z * w
                        p_m = np.array([x_c, y_c, z_px], dtype=np.float32) * meters_per_pixel
                        pts[name] = p_m

                    cur_vec = np.stack([pts[n] for n in POINT_ORDER], axis=0)
                    prev = prev_points[label]
                    cur_vec = ema(prev, cur_vec, EMA_ALPHA)
                    prev_points[label] = cur_vec
                    for i, n in enumerate(POINT_ORDER):
                        pts[n] = cur_vec[i]

                    metrics = compute_metrics(pts)
                    ang_ok, cos_arr = compute_grab_angle_metrics(pts, metrics["tip_names"])

                    t_now = time.time()
                    hc = hold_counters[label]

                    # 抓取：距离 or 角度
                    cond_dist = (metrics["grab_avg"] < GRAB_ENTER_AVG_TIP_DIST_M)
                    near_count = sum((np.linalg.norm(pts[n] - metrics["palm"]) < GRAB_ENTER_AVG_TIP_DIST_M) for n in metrics["tip_names"])
                    cond_near = (near_count >= GRAB_MIN_TIPS_NEAR)
                    cond_ang = (ang_ok >= GRAB_MIN_TIPS_ANG_OK)
                    is_grab_now = (cond_ang or (cond_dist and cond_near))

                    # 抬手
                    is_raise_now = (metrics["raise_y"] > RAISE_Y_THRESHOLD_M) and (metrics["spread"] > RAISE_SPREAD_MIN_M) and (metrics["index_dir_y"] > 0.0)

                    hc["grab"] = hc["grab"] + 1 if is_grab_now else 0
                    hc["raise"] = hc["raise"] + 1 if is_raise_now else 0

                    if hc["grab"] >= GRAB_HOLD_FRAMES and (t_now - last_grab_time[label]) > GRAB_COOLDOWN_S:
                        send_event(net, label, "grab")
                        last_grab_time[label] = t_now
                        hc["grab"] = 0
                    else:
                        if hc["raise"] >= RAISE_HOLD_FRAMES and (t_now - last_raise_time[label]) > RAISE_COOLDOWN_S:
                            send_event(net, label, "raise")
                            last_raise_time[label] = t_now
                            hc["raise"] = 0
                        else:
                            reversals, amplitude = update_wave(label, t_now, metrics["index_tip"][0])
                            if (t_now - last_wave_time[label]) > WAVE_COOLDOWN_S:
                                if (metrics["spread"] > max(RAISE_SPREAD_MIN_M, 0.055)) and (hc["grab"] == 0) and (hc["raise"] == 0):
                                    if reversals >= WAVE_MIN_REVERSALS and amplitude >= WAVE_MIN_AMPLITUDE_M:
                                        send_event(net, label, "wave")
                                        last_wave_time[label] = t_now
                                        wave_history[label].clear()

                    if DEBUG_PRINT and (frame_count % 10 == 0):
                        reversals, amplitude = update_wave(label, t_now, metrics["index_tip"][0])
                        print(f"[DBG] {label:<5s} spread={metrics['spread']:.3f}m avg={metrics['grab_avg']:.3f}m "
                              f"near={near_count} angOK={ang_ok} "
                              f"raise_y={metrics['raise_y']:+.3f}m idx_dir_y={metrics['index_dir_y']:+.3f} "
                              f"rev/amp={reversals}/{amplitude:.3f} "
                              f"cos={np.array2string(cos_arr, precision=2, floatmode='fixed')}")

                    if PRINT_JOINTS_EVERY_N > 0 and (frame_count % PRINT_JOINTS_EVERY_N == 0):
                        for n in ["Wrist", "ThumbTip", "IndexMCP", "IndexTip", "MiddleTip", "RingTip", "PinkyTip"]:
                            x, y, z = pts[n]
                            print(f"  {n:12s} -> x={x:+.3f} y={y:+.3f} z={z:+.3f} (m)")

            info_text = f"{TRANSPORT.upper()} {'%s:%d'%(HOST,PORT) if TRANSPORT!='ws' else WS_URL} | ESC to quit | MIRROR={MIRROR_PREVIEW}"
            cv2.putText(frame, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 0), 2)
            cv2.imshow("Hand Gestures (MediaPipe)", frame)
            if cv2.waitKey(1) & 0xFF == 27:
                break

    finally:
        cap.release()
        cv2.destroyAllWindows()
        hands.close()
        net.close()

if __name__ == "__main__":
    main()